import os, time
import pandas as pd
import numpy as np
# from tqdm import tqdm
import cv2
from pandarallel import pandarallel
pandarallel.initialize()


def segImage(I,S,*args):
    debug = 0
    varargin = args
    boundaryWidth = 1
    boundaryColor = 'red'
    if debug:
        t0 = time.time()
    # Parse the optional inputs.
    if ( len(varargin) % 2 != 0 ):
            print('Extra Parameters passed to the function must be passed in pairs.')

    parameterCount = int(len(varargin)/2)

    for parameterIndex in range(parameterCount):
            parameterName = varargin[parameterIndex*2 - 1]
            parameterValue = varargin[parameterIndex*2]
            switcher = str.lower(parameterName)
            if switcher ==  'boundarycolor':
                boundaryColor = parameterValue
            elif switcher == 'boundarywidth':
                boundaryWidth = parameterValue
            else:
                print('Invalid parameter')
    
    # S = np.array(S, dtype = np.int16)
    S = np.int16(S)
    #S = im2db(S)
    [cy,cx] = np.gradient(S)
    if debug:
        t1 = time.time()
        print('gradient take time = {}'.format(t1 - t0))
    ccc = np.where((abs(cx)+abs(cy)) >0, 1, 0)
    ccc1 = np.where(ccc > 0, 0, 1)

    # ccc = ccc.astype(int)
    # ccc1 = ccc1.astype(int)
    ccc = np.uint8(ccc)
    ccc1 = np.uint8(ccc1)
    if boundaryWidth>1:
        boundaryWidth = np.ceil(boundaryWidth)
        dilateWindow = np.ones(boundaryWidth, boundaryWidth)
        ccc = cv2.dilate(ccc, dilateWindow, iterations=1) #imdilate(ccc,dilateWindow)
    elif boundaryWidth<1:
        print('boundaryWidth has been reset to 1.')

    if boundaryColor == 'red':
        I[:,:,0] = np.maximum(I[:,:,0],ccc)
        I[:,:,1] = np.minimum(I[:,:,1],ccc1)
        I[:,:,2] = np.minimum(I[:,:,2],ccc1)
    elif boundaryColor== 'black':
        I[:,:,0] = np.minimum(I[:,:,0],ccc1)
        I[:,:,1] = np.minimum(I[:,:,1],ccc1)
        I[:,:,2] = np.minimum(I[:,:,2],ccc1)
    else:
        print('Does not recognize boundaryColor other than red and black')
        
        I[:,:,0] = np.maximum(I[:,:,0],ccc1)
        I[:,:,1] = np.maximum(I[:,:,1],ccc1)
        I[:,:,2] = np.maximum(I[:,:,2],ccc1)
    if debug:
        t2 = time.time()
        print('segImg take time = {}'.format(t2 - t1))

    return I


img_path = '/opt/data/private/project/adc_T9/50UD_05300/selected_train_data/Images'
mask_path = '/opt/data/private/project/adc_T9/50UD_05300/selected_train_data/seg_pred_mask'
dst_path = '/opt/data/private/project/adc_T9/50UD_05300/selected_train_data/Images_add_mask2'
Wm = 2044
Hm = 2044

df = pd.DataFrame()
codes = []
imgs = []
for code in os.listdir(img_path):
    tmp = os.listdir(os.path.join(img_path, code))
    codes.extend([code] * len(tmp))
    imgs.extend(tmp[:])

# df['root_path'] = [img_path] * len(codes)
df['code'] = codes
df['image'] = imgs

def funcx(code, img_name):
    mask_name = img_name[:-3] + 'png'
    img = cv2.imread(os.path.join(img_path, code, img_name), 1)
    mask = cv2.imread(os.path.join(mask_path, code, mask_name), 0)
    mask = cv2.resize(mask, (Wm, Hm), interpolation=cv2.INTER_NEAREST)
    masked_img = segImage(img/255, mask)
    os.makedirs(os.path.join(dst_path, code), exist_ok=True)
    cv2.imwrite(os.path.join(dst_path, code, mask_name), masked_img*255)

df.parallel_apply(lambda x: funcx(x['code'], x['image']), axis=1) 
